from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.backends import default_backend
import base64
# 生成私钥
private_key = rsa.generate_private_key(
    public_exponent=65537,
    key_size=2048,
    backend=default_backend()
)

# 获取公钥
public_key = private_key.public_key()

# 序列化私钥为PEM格式
pem_private = private_key.private_bytes(
    encoding=serialization.Encoding.PEM,
    format=serialization.PrivateFormat.PKCS8,  # 注意这里是 PrivateFormat 不是 PrivateFormat
    encryption_algorithm=serialization.NoEncryption()
).decode('utf-8')

# 序列化公钥为PEM格式
pem_public = public_key.public_bytes(
    encoding=serialization.Encoding.PEM,
    format=serialization.PublicFormat.SubjectPublicKeyInfo
).decode('utf-8')

from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives import hashes

def verify_key_pair_with_signature(public_key, private_key):
    # 测试消息
    message = b"Test message for key verification"
    
    # 使用私钥签名
    private_key = serialization.load_pem_private_key(
            private_key.encode(),
            password=None,
            backend=default_backend()
    )
    signature = private_key.sign(
        message,
        padding.PSS(
            mgf=padding.MGF1(hashes.SHA256()),
            salt_length=padding.PSS.MAX_LENGTH
        ),
        hashes.SHA256()
    )
    signature= base64.b64encode(signature).decode('utf-8')
    # 使用公钥验证签名
    public_key = serialization.load_pem_public_key(
        public_key.encode(),
        backend=default_backend()
    )
    try:
        public_key.verify(
            base64.b64decode(signature),
            message,
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()),
                salt_length=padding.PSS.MAX_LENGTH
            ),
            hashes.SHA256()
        )
        return True
    except:
        return False

print(verify_key_pair_with_signature(pem_public, pem_private))

public_key = '''-----BEGIN PUBLIC KEY-----
MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAsO9JqRo02Zq5CIv1/QFD
dPKy+Mn+OitNLNfe7BMtIUj/YuesWmW6pc9YAxv9MZCDJdiWcEpQR5Th4XKHtPJR
PWo3w4FJsXxF+3zixhc9Im52GxTTztDhKSgBbmFIIDqM3JBDAzX6da+4LONyltOp
aCNJRh11+ekTVizkSNdGhuPJHueseqs54rUKyTIckeWtGzJtbDggTumNFLUMqfHN
E4WLpXNFKX/w1u2d3iSvwFWsTwO80xr4lYcj9OTpOlABiKAG1UjwS9h2LwcxMZml
qScJxwy28BxOHO2ItSiwrLLOPmcHsUSjBFlwPjdk/yBvTVeNG04HgJoO14bAn0uF
nQIDAQAB
-----END PUBLIC KEY-----
'''
public_key = serialization.load_pem_public_key(
    public_key.encode(),
    backend=default_backend()
)
signature="KNhgkNHq3CZaM6Zstp8KTR3M12JrKuL5Ysg5Z2+89bKGPg8xNiV44cpuCwmbmP1Ip4RY/Xehp+AZRjOu7lgeUWfx4lN834/MUjs3JheFQpRWpVcxsBflVlCmZ4Nm+21An1IVo1GXDhhVXGH8VK2GsuAcW30K/9v85bG+Ol2saSJT8BzAP9dgYFgh96rYDqjrLzhyvifCpoiLpbwmsi0kgcDFzfhcvmMkwteR+Xhm2oovBCrkOG1fR+PVM5W9Ky5HDG8A5p2RI8o+KJ+8swxezgrKR8W6mI4gVghjdWQEoj8rNRuX75tOOCr70V+de3KWMLRKReeq3IFzRdDFOXGw5Q=="
message = b"Test message for key verificatio"
public_key.verify(
    base64.b64decode(signature),
    message,
    padding.PSS(
        mgf=padding.MGF1(hashes.SHA256()),
        salt_length=padding.PSS.MAX_LENGTH
    ),
    hashes.SHA256()
)


# public_key='''
# -----BEGIN PUBLIC KEY-----
# MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA4SIydazIgXQ8SYOdEPvB
# sxlT7945MWc+L47WOk9NLauD2XYBfIjs3gg3Wx/k3m1e3gfKlPiXtecsWSsS2m8Z
# ViMSZDHlm3NkDKpm4aJbFYCaH2vRFrlUialY7/dVZV6+C92GPTl0SNOv4GNBhi/e
# 0+YG38VyQZjGZu8UeA4PPEzN3A4BQMUzhMq4V90JRJhZvGUXRylez9NbOBezE6D+
# Ijub7WqTKg5XIFU4NeRD8S9MgO3jrJVpajuYZ3QefLtJOPbCG0lN0OoB2xwAGrO6
# rD88g1Hcs1f8ekWKz14RHoIwFwCgBSEmF3wYO6pmoESYACehR2DB+TtvwzJgvGB5
# ZQIDAQAB
# -----END PUBLIC KEY-----
# '''
# public_key = serialization.load_pem_public_key(
#     public_key.encode(),
#     backend=default_backend()
# )

# signature='''2Tg+KiQElUfT6oQssgMa+fJInw5HfCiE1ez57Qr2LUO1JNdDN30Q18VAJzioI+bygYCQOICX/VMsgwj7a6VSW1MPArEh1WppvYBYNuoT1AWx6RNOOUy9WlBLq8kQuMFg7RlhIpFnM4qv+PoDSGbIF3iwZW7k/VR1ywgK1ICPaa2Wf4vVp5DvTUTPFKonK5d05dJf0OtIib2iNo+SWd2HMVCFNp3ZrQBuPWodtdVKEFzAp3HTVle9gxkOax5wl+wxqTfW0LbDbh2ZY5sf4mrhOEjnAOMNmFX+/I46C4ifODKaR0CJgbFsWGPgADky6fEArf6LerejfEFLKPAa4UAsoA=='''

# message = b"Test message for key verification"
# print(public_key.verify(
#     base64.b64decode(signature),
#     message,
#     padding.PSS(
#         mgf=padding.MGF1(hashes.SHA256()),
#         salt_length=padding.PSS.MAX_LENGTH
#     ),
#     hashes.SHA256()
# ))
